// Copyright 2022 Chen Jun

#include "hnurm_armor/tracker.hpp"
#include <angles/angles.h>
// #include "easylogging++.h"

// STD
#include <cfloat>
#include <opencv2/calib3d.hpp>
#include <rclcpp/rclcpp.hpp>
#include <string>

#define CLAMP(x, a, b) (MIN(MAX(x, a), b))

namespace hnurm
{
    Tracker::Tracker(rclcpp::Node::SharedPtr node)
    {
        tracker_state = LOST;
        tracked_id = std::string("");
        target_state = Eigen::VectorXd::Zero(9);

        max_match_distance_ = node->declare_parameter("processor.tracker.max_match_distance", 0.5);
        max_match_yaw_diff_ = node->declare_parameter("processor.tracker.max_match_yaw_diff", 0.5);
        tracking_threshold_ = node->declare_parameter("processor.tracker.tracking_threshold", 0.0);
        lost_threshold_ = node->declare_parameter("processor.tracker.lost_threshold", 100.0);

        measurement = Eigen::Vector4d::Zero();
    }

    Tracker::Tracker(double max_match_distance, int tracking_threshold, int lost_threshold)
        : tracker_state(LOST),
          tracked_id(std::string("")),
          target_state(Eigen::VectorXd::Zero(9)),
          max_match_distance_(max_match_distance),
          tracking_threshold_(tracking_threshold),
          lost_threshold_(lost_threshold)
    {
    }

    void Tracker::init(const std::vector<Armor> &armors_msg)
    {
        if (armors_msg.empty())
        {
            return;
        }
        // Simply choose the armor that is closest to image center
        auto min_distance = DBL_MAX;
        tracked_armor = armors_msg[0];
        for (const auto &armor: armors_msg)
        {
            if (armor.distance_to_image_center < min_distance)
            {
                min_distance = armor.distance_to_image_center;
                tracked_armor = armor;
            }
        }

        initEKF(tracked_armor);

        tracked_id = tracked_armor.number;
        tracker_state = DETECTING;

        updateArmorsNum(tracked_armor);
    }

    void Tracker::update(const std::vector<Armor> &armors_msg)
    {
        // KF predict
        Eigen::VectorXd ekf_prediction = ekf.predict();
        // RCLCPP_DEBUG(rclcpp::get_logger("armor_processor"), "EKF predict");

        bool matched = false;
        // Use KF prediction as default target state if no matched armor is found
        target_state = ekf_prediction;

        if (!armors_msg.empty())
        {
            auto min_position_diff = DBL_MAX;
            auto yaw_diff = DBL_MAX;

            auto predicted_position = getArmorPositionFromState(ekf_prediction);
            for (const auto &armor: armors_msg)
            {
                // 仅考虑相同ID的装甲板
                if (armor.number == tracked_id)
                {
                    auto p = armor.position;
                    // 当前装甲板位置与跟踪装甲板的预测位置的差异
                    double position_diff = (predicted_position - p).norm();
                    if (position_diff < min_position_diff)
                    {
                        min_position_diff = position_diff;
                        // 计算当前装甲板的姿态与EKF预测的姿态之间的差异
                        yaw_diff = abs(orientationToYaw(armor.rotation) - ekf_prediction(6));
                        tracked_armor = armor;
                    }
                }
            }
            // 保存差异信息
            info_position_diff = min_position_diff;
            info_yaw_diff = yaw_diff;

            if (min_position_diff < max_match_distance_ && yaw_diff < max_match_yaw_diff_)
            {
                // 找到匹配的装甲板
                matched = true;
                auto p = tracked_armor.position;
                // 更新EKF
                double measured_yaw = orientationToYaw(tracked_armor.rotation);
                Eigen::Vector4d z(p(0), p(1), p(2), measured_yaw);
                target_state = ekf.update(z);
            }
            else if (yaw_diff > max_match_yaw_diff_) // 此时发生了装甲板跳变
            {
                // 处理装甲板跳变情况
                handleArmorJump(tracked_armor);
                if (tracker_state == TRACKING)
                {
                    tracker_state = JUMP;
                }
            }
            else
            {
                // 装甲板未匹配，追踪状态设置为丢失
                tracker_state = LOST;
            }
        }

        // Prevent radius from spreading
        target_state(8) = CLAMP(target_state(8), 0.14, 0.33);
        another_r = CLAMP(another_r, 0.14, 0.33);
        // Tracking state machine
        if (tracker_state == DETECTING)
        {
            if (matched)
            {
                detect_count_++;
                if (detect_count_ > tracking_threshold_)
                {
                    detect_count_ = 0;
                    tracker_state = TRACKING;
                }
            }
            else
            {
                detect_count_ = 0;
                tracker_state = LOST;
            }
        }
        else if (tracker_state == TRACKING || tracker_state == JUMP)
        {
            if (!matched && tracker_state != JUMP)
            {
                tracker_state = TEMP_LOST;
                lost_count_++;
            }
            if (matched && tracker_state == JUMP)
            {
                tracker_state = TRACKING;
            }
            if (!matched && tracker_state == JUMP)
            {
                lost_count_++;
                if (lost_count_ > lost_threshold_)
                {
                    lost_count_ = 0;
                    tracker_state = LOST;
                }
            }
        }
        else if (tracker_state == TEMP_LOST)
        {
            if (!matched)
            {
                lost_count_++;
                if (lost_count_ > lost_threshold_)
                {
                    lost_count_ = 0;
                    tracker_state = LOST;
                }
            }
            else
            {
                tracker_state = TRACKING;
                lost_count_ = 0;
            }
        }
    }

    void Tracker::initEKF(const Armor &a)
    {
        double xa = a.position(0);
        double ya = a.position(1);
        double za = a.position(2);
        last_yaw_ = 0;
        // todo
        double yaw = orientationToYaw(a.rotation);

        // Set initial position at 0.2m behind the target
        target_state = Eigen::VectorXd::Zero(9);
        double r = 0.20;
        double xc = xa + r * cos(yaw);
        double yc = ya + r * sin(yaw);
        dz = 0, another_r = r;
        target_state << xc, 0, yc, 0, za, 0, yaw, 0, r;

        ekf.setState(target_state);
        // RCLCPP_DEBUG(rclcpp::get_logger("armor_processor"), "Init EKF!");
    }

    void Tracker::handleArmorJump(const Armor &a)
    {
        // todo
        double yaw = orientationToYaw(a.rotation);
        target_state(6) = yaw;

        updateArmorsNum(a);

        if (tracked_armors_num == ArmorsNum::NORMAL_4)
        {
            dz = target_state(4) - a.position(2);
            target_state(4) = a.position(2);
            std::swap(target_state(8), another_r);
        }

        auto p = a.position;

        Eigen::Vector3d infer_p = getArmorPositionFromState(target_state);

        if ((p - infer_p).norm() > max_match_distance_)
        {
            double r = target_state(8);
            target_state(0) = p(0) + r * cos(yaw);  // xc
            target_state(1) = 0;                    // vxc
            target_state(2) = p(1) + r * sin(yaw);  // yc
            target_state(3) = 0;                    // vyc
            target_state(4) = p(2);                 // za
            target_state(5) = 0;                    // vza
            // CLOG(WARNING, "process") << "Reset state!";
            RCLCPP_WARN(rclcpp::get_logger("armor_processor"), "Reset state!");
        }

        ekf.setState(target_state);
    }

    double Tracker::orientationToYaw(const Eigen::Matrix3d &R)
    {
        double sy = sqrt(R(0, 0) * R(0, 0) + R(1, 0) * R(1, 0));
        Eigen::Vector3d eulerAngles;
        bool singular = sy < 1e-6; // If

        if (!singular)
        {
            eulerAngles[0] = atan2(R(2, 1), R(2, 2)); // X-axis rotation
            eulerAngles[1] = atan2(-R(2, 0), sy); // Y-axis rotation
            eulerAngles[2] = atan2(R(1, 0), R(0, 0)); // Z-axis rotation
        }
        else
        {
            eulerAngles[0] = atan2(-R(1, 2), R(1, 1)); // X-axis rotation
            eulerAngles[1] = atan2(-R(2, 0), sy); // Y-axis rotation
            eulerAngles[2] = 0; // Z-axis rotation
        }
        double yaw = eulerAngles[2] + CV_PI / 2;
        armor_pitch = eulerAngles[0] / CV_PI * 180;
        armor_yaw = eulerAngles[2] / CV_PI * 180;
//    double yaw = atan2(R(1, 1), R(0, 1));
//    armor_pitch=atan2(R(2,2),R(1,2))/CV_PI*180;
//    armor_yaw=yaw/CV_PI*180;
        yaw = last_yaw_ + angles::shortest_angular_distance(last_yaw_, yaw);
        last_yaw_ = yaw;
        return yaw;
    }

    Eigen::Vector3d Tracker::getArmorPositionFromState(const Eigen::VectorXd &x)
    {
        // Calculate predicted position of the current armor
        double xc = x(0), yc = x(2), zc = x(4);
        double yaw = x(6), r = x(8);
        double xa = xc - r * cos(yaw);
        double ya = yc - r * sin(yaw);
        return {xa, ya, zc};
    }

    void Tracker::updateArmorsNum(const Armor &armor)
    {
        if (armor.armor_type == LARGE && (tracked_id == "3" || tracked_id == "4" || tracked_id == "5"))
        {
            tracked_armors_num = ArmorsNum::BALANCE_2;
        }
        else if (tracked_id == "outpost")
        {
            tracked_armors_num = ArmorsNum::OUTPOST_3;
        }
        else if (tracked_id == "base")
        {
            tracked_armors_num = ArmorsNum::BASE_1;
        }
        else
        {
            tracked_armors_num = ArmorsNum::NORMAL_4;
        }
    }

}  // namespace hnurm
